package com.alan.test;

import io.netty.buffer.EmptyByteBuf;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Objects;

/**
 * 请求、响应拦截器
 *
 * @author tangtong
 * @date 2020/12/18
 */
@Slf4j
@Component
public class ReqResGlobalFilter implements GlobalFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 未开启打印日志，直接下一个过滤器

        long start = System.currentTimeMillis();

        ServerHttpRequest request = exchange.getRequest();
        MediaType mediaType = request.getHeaders().getContentType();

        // 构造操作日志对象
        final OperateLog operateLog = new OperateLog();
//        operateLog.setTraceId(tracer.currentSpan().context().traceIdString());
        operateLog.setUri(request.getURI().getPath());
        operateLog.setMethod(request.getMethodValue());
        operateLog.setContentType(String.valueOf(mediaType));
        operateLog.setQueryString(request.getURI().getQuery());

        // 装饰response
        ServerHttpResponseDecorator responseDecorator = decorateResponse(exchange, operateLog);

        // 装饰request
        return logRequest(exchange.mutate().response(responseDecorator).build(), chain, operateLog).doFinally(s -> {
            long end = System.currentTimeMillis();

            operateLog.setElapsedTime(end - start);

            // 记录日志的操作
            log.info("operate log = {}", operateLog);
        });
    }

    private Mono<Void> logRequest(ServerWebExchange exchange, GatewayFilterChain chain, OperateLog operateLog) {
        MediaType mediaType = exchange.getRequest().getHeaders().getContentType();
        boolean isJson = Objects.equals(mediaType, MediaType.APPLICATION_JSON) || Objects.equals(mediaType, MediaType.APPLICATION_JSON_UTF8);
        boolean isFormUrlEncoded = Objects.equals(mediaType, MediaType.APPLICATION_FORM_URLENCODED);
        boolean isLogin = exchange.getRequest().getURI().getPath().contains("login");

        // 登录接口，不记录请求体（防止用户名密码泄漏）
        if (!isLogin) {
            // json或者x-www-form-urlencoded形式的参数才记录，form-data是上传文件的，不记录
            if (isJson || isFormUrlEncoded) {
                // 构建DataBuffer的工厂
                NettyDataBufferFactory factory = (NettyDataBufferFactory) exchange.getResponse().bufferFactory();

                // 聚合，等所有RequestBody都接收到了才往下进行
                return DataBufferUtils.join(exchange.getRequest().getBody())
                        // 如果body为空则构造一个空的
                        .defaultIfEmpty(factory.wrap(new EmptyByteBuf(factory.getByteBufAllocator())))
                        // 封装一个RequestDecorator
                        .map(dataBuffer -> decorateRequest(exchange, dataBuffer, factory, operateLog))
                        // 如果为空，还使用原来的request
                        .switchIfEmpty(Mono.just(exchange.getRequest()))
                        // 构造返回值
                        .flatMap(req -> {
                            // 如果没被装饰过，直接返回
                            if (req == exchange.getRequest()) {
                                return chain.filter(exchange);
                            }
                            // 重新构建一个exchange对象
                            return chain.filter(exchange.mutate().request(req).build());
                        });
            }
        }
        return chain.filter(exchange);
    }

    private ServerHttpRequest decorateRequest(ServerWebExchange exchange, DataBuffer dataBuf, NettyDataBufferFactory factory, OperateLog operateLog) {

        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public Flux<DataBuffer> getBody() {
                // 解析DataBuffer
                byte[] buffer = new byte[dataBuf.readableByteCount()];
                dataBuf.read(buffer);
                String content = new String(buffer, StandardCharsets.UTF_8);

                // 释放，防止直接内存溢出
                DataBufferUtils.release(dataBuf);

                // 记录到日志中
                operateLog.setRequest(content);

                // 重新构建一个DataBuffer往下传
                return Mono.just(factory.wrap(content.getBytes(StandardCharsets.UTF_8))).flux();
            }
        };
    }

    private ServerHttpResponseDecorator decorateResponse(ServerWebExchange exchange, final OperateLog operateLog) {
        // 构建DataBuffer的工厂
        NettyDataBufferFactory factory = (NettyDataBufferFactory) exchange.getResponse().bufferFactory();

        return new ServerHttpResponseDecorator(exchange.getResponse()) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

                return super.writeWith(DataBufferUtils.join(body)
                        .defaultIfEmpty(factory.wrap(new EmptyByteBuf(factory.getByteBufAllocator())))
                        .map(dataBuf -> {
                            MediaType mediaType = getHeaders().getContentType();
                            boolean isJson = Objects.equals(mediaType, MediaType.APPLICATION_JSON)
                                    || Objects.equals(mediaType, MediaType.APPLICATION_JSON_UTF8);

                            // 返回值类型为json时，才拦截
                            if (isJson) {
                                int readPosition = dataBuf.readPosition();
                                // 解析DataBuffer
                                byte[] buffer = new byte[dataBuf.readableByteCount()];
                                dataBuf.read(buffer);
                                String content = new String(buffer, StandardCharsets.UTF_8);

                                // 释放，防止直接内存溢出
//                                DataBufferUtils.release(dataBuf);

                                // 记录到日志中
                                operateLog.setResponse(content);

                                dataBuf.readPosition(readPosition);
                                // 重建一个DataBuffer往下传
                                return dataBuf;
//                                return bufferFactory().wrap(content.getBytes(StandardCharsets.UTF_8));
                            }
                            return dataBuf;
                        }));

            }
        };
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 2;
    }
}
